#!/usr/bin/env julia

"""
Execute the MWRI regional GSI workflow with the current Julia implementation.

Workflow summary:
  1. Read the shared WRF background (`wrfinput_d01`).
  2. Flatten the fields into the Julia state ordering (u, v, T, q, ps).
  3. Decode PrepBUFR winds, map them onto the regional grid, and perform either a
     diagonal 3D-Var update or a NAM-style CG solve using the operational B operator.
  4. Sample ATMS/MHS/CrIS radiance BUFR files to verify decode coverage.
  5. Emit analysis metrics plus background/analysis/increment binary dumps.

The script provides scaffolding for future parity checks. Radiance assimilation
and full CRTM coupling remain follow-up tasks once the wind/NAM workflow is stable.
"""

using Dates
using LinearAlgebra
using Printf
using Statistics
using JSON3
using Random

using GSICoreAnalysis
using GSICoreAnalysis.DataIO
using GSICoreAnalysis.DataIO.WindBUFRReader
using GSICoreAnalysis.GridOperations: get_coordinates
using GSICoreAnalysis.DataIO.RadianceBUFRReader: RadianceBUFRConfig, read_radiance_bufr
using GSICoreAnalysis.BackgroundError
using GSICoreAnalysis.ControlVectors: ControlVector

const DEFAULT_BACKGROUND =
    "/home/docker/GSI/comGSIv3.7_EnKFv1.3-mwri/test/WRF-4.3.3/wrfinput_d01"
const DEFAULT_PREPBUFR =
    "/home/docker/GSI/comGSIv3.7_EnKFv1.3-mwri/test/satellite_gdas/gdas.t00z.prepbufr.nr"
const DEFAULT_OUTPUT_DIR = "results/mwri_regional_case_dev"

const DEFAULT_RADIANCE_FILES = [
    ("ATMS", "/home/docker/GSI/comGSIv3.7_EnKFv1.3-mwri/test/satellite_gdas/gdas.t00z.atms.tm00.bufr_d"),
    ("MHS",  "/home/docker/GSI/comGSIv3.7_EnKFv1.3-mwri/test/satellite_gdas/gdas.t00z.1bmhs.tm00.bufr_d"),
    ("CrIS", "/home/docker/GSI/comGSIv3.7_EnKFv1.3-mwri/test/satellite_gdas/gdas.t00z.crisf4.tm00.bufr_d")
]

const SCENARIO_DEFAULTS = Dict(
    "mwri" => (
        background = DEFAULT_BACKGROUND,
        prepbufr = DEFAULT_PREPBUFR,
        output = DEFAULT_OUTPUT_DIR,
        radiance_files = DEFAULT_RADIANCE_FILES,
        radiance_sample = 10_000,
        analysis_mode = "diag",
        bufr_table = nothing,
        bufr_backend = :bufrlib,
    ),
    "tutorial" => (
        background = "/home/docker/comgsi/tutorial/case_data/2018081212/bkg/wrfinput_d01.mem0001",
        prepbufr = "/home/docker/comgsi/tutorial/case_data/2018081212/obs/rap.t12z.prepbufr.tm00",
        output = "results/tutorial_regional_case_dev",
        radiance_files = Tuple{String,String}[],
        radiance_sample = 0,
        analysis_mode = "diag",
        bufr_table = "/home/docker/comgsi/tutorial/comGSIv3.7_EnKFv1.3/scripts/fix/prepobs_prep.bufrtable",
        bufr_backend = :bufrlib,
    )
)

const BACKGROUND_WIND_STD = Dict(:u => 3.5, :v => 3.5) # m s⁻¹ diagonal-B proxy

# Platform-specific observation error standard deviations (m/s)
# Based on operational GSI error assignments
const PLATFORM_ERROR_STDS = Dict(
    120 => 2.0,    # RAWINSONDE (RAOB) - high quality
    220 => 2.5,    # RAWINSONDE dropwindsonde
    221 => 2.5,    # PIBAL
    130 => 2.5,    # AIRCRAFT (AIREP, ACARS)
    131 => 2.5,    # AIRCRAFT
    133 => 2.5,    # AIRCRAFT MDCRS
    180 => 3.0,    # SURFACE SYNOP
    181 => 4.0,    # SHIP
    182 => 4.0,    # BUOY (drifting)
    183 => 4.0,    # BUOY (moored)
    187 => 4.0,    # METAR
    242 => 3.5,    # SATWIND (goes)
    243 => 3.5,    # SATWIND (modis/terra)
    244 => 3.5,    # SATWIND (modis/aqua)
    245 => 3.5,    # SATWIND
    246 => 3.5,    # SATWIND
    250 => 3.5,    # SATWIND (various satellites)
    251 => 3.5,    # SATWIND
    252 => 3.5,    # SATWIND
    253 => 3.5,    # SATWIND
    254 => 3.5,    # SATWIND
    257 => 3.5,    # SATWIND
    258 => 3.5,    # SATWIND
    259 => 3.5,    # SATWIND
)

const DEFAULT_WIND_ERROR_STD = 4.5  # Default for unknown platform types

struct WindObservationPayload
    values::Vector{Float64}
    errors::Vector{Float64}
    state_indices::Vector{Int}
    components::Vector{Symbol}
    platforms::Vector{Int}
    counts::Dict{Symbol,Int}
end

struct NAMOperators
    config::AnalysisConfig{Float64}
    B::BackgroundError.OperationalBMatrix{Float64}
    cv_workspace::ControlVector{Float64}
end

function ensure_eccodes_overlay!()
    GSICoreAnalysis.DataIO.BUFRReader._configure_eccodes_environment!()
    GSICoreAnalysis.DataIO.RadianceBUFRReader._configure_eccodes_environment!()
end

function parse_cli_args()
    scenario = "mwri"
    for arg in ARGS
        if startswith(arg, "--scenario=")
            scenario = lowercase(String(split(arg, "=", limit=2)[2]))
        end
    end

    defaults = get(SCENARIO_DEFAULTS, scenario, nothing)
    defaults === nothing && error("Unknown scenario '$scenario'. Supported scenarios: $(join(keys(SCENARIO_DEFAULTS), ", ")).")

    background = defaults.background
    prepbufr = defaults.prepbufr
    output_dir = defaults.output
    radiance_sample = defaults.radiance_sample
    analysis_mode = defaults.analysis_mode
    bufr_table_path = defaults.bufr_table
    bufr_backend = defaults.bufr_backend
    radiance_files = copy(defaults.radiance_files)

    # NAM solver tuning parameters
    nam_max_iter = 40
    nam_tol = 5e-3
    nam_preconditioner = "diag"
    nam_use_balance = false
    nam_adaptive_tol = true

    # Observation processing parameters
    obs_thinning_distance = 0.0  # km (0 = no thinning)
    obs_max_count = 0  # Maximum observations (0 = unlimited)
    obs_error_floor = 2.0  # Minimum observation error (m/s)
    obs_use_platform_errors = true  # Use platform-specific error assignment

    # Background error tuning
    background_wind_std = 5.0  # Background error std dev for winds (m/s) - INCREASED from 3.5 for better J_b/J_o balance

    for arg in ARGS
        if startswith(arg, "--scenario=")
            # already handled in first pass
            continue
        elseif startswith(arg, "--background=")
            background = String(split(arg, "=", limit=2)[2])
        elseif startswith(arg, "--prepbufr=")
            prepbufr = String(split(arg, "=", limit=2)[2])
        elseif startswith(arg, "--output=")
            output_dir = String(split(arg, "=", limit=2)[2])
        elseif startswith(arg, "--radiance-sample=")
            radiance_sample = parse(Int, split(arg, "=", limit=2)[2])
        elseif startswith(arg, "--analysis-mode=")
            analysis_mode = lowercase(String(split(arg, "=", limit=2)[2]))
        elseif startswith(arg, "--bufr-table=")
            bufr_table_path = String(split(arg, "=", limit=2)[2])
        elseif startswith(arg, "--bufr-backend=")
            backend_str = lowercase(String(split(arg, "=", limit=2)[2]))
            backend_sym = Symbol(backend_str)
            backend_sym ∈ (:bufrlib, :eccodes, :auto) || error("Unsupported bufr-backend '$backend_str'. Use bufrlib|eccodes|auto.")
            bufr_backend = backend_sym
        elseif startswith(arg, "--nam-max-iter=")
            nam_max_iter = parse(Int, split(arg, "=", limit=2)[2])
        elseif startswith(arg, "--nam-tol=")
            nam_tol = parse(Float64, split(arg, "=", limit=2)[2])
        elseif startswith(arg, "--nam-preconditioner=")
            nam_preconditioner = lowercase(String(split(arg, "=", limit=2)[2]))
        elseif arg == "--nam-use-balance"
            nam_use_balance = true
        elseif arg == "--nam-no-adaptive-tol"
            nam_adaptive_tol = false
        elseif startswith(arg, "--obs-thinning-distance=")
            obs_thinning_distance = parse(Float64, split(arg, "=", limit=2)[2])
        elseif startswith(arg, "--obs-max-count=")
            obs_max_count = parse(Int, split(arg, "=", limit=2)[2])
        elseif startswith(arg, "--obs-error-floor=")
            obs_error_floor = parse(Float64, split(arg, "=", limit=2)[2])
        elseif arg == "--obs-no-platform-errors"
            obs_use_platform_errors = false
        elseif startswith(arg, "--background-wind-std=")
            background_wind_std = parse(Float64, split(arg, "=", limit=2)[2])
        elseif arg == "--help" || arg == "-h"
            println("Usage: run_mwri_regional_case.jl [OPTIONS]")
            println("\nGeneral Options:")
            println("  --scenario=tutorial|mwri    Select preset configuration (default: mwri)")
            println("  --background=PATH           WRF background file")
            println("  --prepbufr=PATH             PrepBUFR observation file")
            println("  --output=DIR                Output directory")
            println("  --analysis-mode=diag|nam    Analysis method (default: diag)")
            println("  --radiance-sample=N         Max radiance obs to sample (default: 10000)")
            println("  --bufr-table=FILE           Optional PrepBUFR table override")
            println("  --bufr-backend=MODE         BUFR backend: bufrlib | eccodes | auto")
            println("\nNAM Solver Tuning Options:")
            println("  --nam-max-iter=N            Maximum CG iterations (default: 40)")
            println("  --nam-tol=TOL               CG residual tolerance (default: 5e-3)")
            println("  --nam-preconditioner=TYPE   Preconditioner: diag|block|none (default: diag)")
            println("  --nam-use-balance           Enable background balance operator")
            println("  --nam-no-adaptive-tol       Disable adaptive RMS convergence (RMS<1.2)")
            println("\nObservation Processing Options:")
            println("  --obs-thinning-distance=KM  Spatial thinning distance in km (default: 0, no thinning)")
            println("  --obs-max-count=N           Maximum observation count (default: 0, unlimited)")
            println("  --obs-error-floor=STD       Minimum observation error in m/s (default: 2.0)")
            println("  --obs-no-platform-errors    Disable platform-specific error assignment")
            println("\nBackground Error Tuning:")
            println("  --background-wind-std=STD   Background error std dev for winds in m/s (default: 3.5)")
            println("\nScenario Defaults:")
            for (key, cfg) in SCENARIO_DEFAULTS
                println("  Scenario: $key")
                println("    background     : $(cfg.background)")
                println("    prepbufr       : $(cfg.prepbufr)")
                println("    output         : $(cfg.output)")
                println("    analysis_mode  : $(cfg.analysis_mode)")
                println("    radiance_sample: $(cfg.radiance_sample)")
                buf_default = cfg.bufr_table === nothing ? "(none)" : String(cfg.bufr_table)
                println("    bufr_table     : $(buf_default)")
            end
            exit(0)
        else
            @warn "Unrecognised argument $arg"
        end
    end

    return (scenario, background, prepbufr, output_dir, radiance_sample, analysis_mode, bufr_table_path,
            bufr_backend, radiance_files, nam_max_iter, nam_tol, nam_preconditioner, nam_use_balance, nam_adaptive_tol,
            obs_thinning_distance, obs_max_count, obs_error_floor, obs_use_platform_errors, background_wind_std)
end

function load_background(path::AbstractString)
    isfile(path) || error("WRF background file not found at $path")
    println("[1/6] Reading WRF background …")
    bg = read_background_fields(path; format=:netcdf, model_type=:wrf)
    grid_cfg = bg.grid.config
    println(@sprintf("  Grid: %d × %d × %d (dx=%.0f m, dy=%.0f m)",
                     grid_cfg.nx, grid_cfg.ny, grid_cfg.nsig,
                     grid_cfg.dx, grid_cfg.dy))
    println(@sprintf("  Background valid time: %s", string(bg.time)))
    return bg
end

function background_to_state_vector(bg)
    grid_cfg = bg.grid.config
    nx, ny, nz = grid_cfg.nx, grid_cfg.ny, grid_cfg.nsig

    u = permutedims(bg.u, (2, 1, 3))
    v = permutedims(bg.v, (2, 1, 3))
    t = permutedims(bg.t, (2, 1, 3))
    q = permutedims(bg.q, (2, 1, 3))
    ps = permutedims(bg.ps, (2, 1))

    n3d = nx * ny * nz
    n2d = nx * ny
    state = Vector{Float64}(undef, 4 * n3d + n2d)

    state[1:n3d] .= vec(u)
    state[(n3d + 1):(2 * n3d)] .= vec(v)
    state[(2 * n3d + 1):(3 * n3d)] .= vec(t)
    state[(3 * n3d + 1):(4 * n3d)] .= vec(q)
    state[(4 * n3d + 1):end] .= vec(ps)

    return state
end

nearest_index_sorted(vec::AbstractVector{<:Real}, val::Real) = nearest_index_sorted(vec, Float64(val))

function nearest_index_sorted(vec::AbstractVector{<:Real}, val::Float64)
    if isnan(val)
        return nothing
    elseif val <= vec[1]
        return 1
    elseif val >= vec[end]
        return length(vec)
    else
        pos = searchsortedfirst(vec, val)
        if pos > length(vec)
            return length(vec)
        elseif pos == 1
            return 1
        else
            return abs(vec[pos] - val) <= abs(val - vec[pos-1]) ? pos : pos - 1
        end
    end
end

function pressure_to_sigma_index(pressure_value::Float64, sigma_axis::AbstractVector{Float64})
    if !isfinite(pressure_value)
        return nothing
    end
    p = pressure_value > 5000 ? pressure_value * 0.01 : pressure_value  # Pa → hPa if needed
    if p <= 0 || p > 1100
        return nothing
    end
    sigma_target = clamp(p / 1000.0, sigma_axis[1], sigma_axis[end])
    return nearest_index_sorted(sigma_axis, sigma_target)
end

function build_wind_payload(prepbufr_path::AbstractString, grid, background_state::Vector{Float64},
                            lon_grid::Matrix{Float64}, lat_grid::Matrix{Float64},
                            sigma_axis::Vector{Float64}, nx::Int, ny::Int, nz::Int;
                            bufr_table_path::Union{Nothing,String}=nothing,
                            bufr_backend::Symbol=:auto,
                            thinning_distance::Float64=0.0,
                            max_count::Int=0,
                            error_floor::Float64=2.0,
                            use_platform_errors::Bool=true)
    println("[3/6] Decoding PrepBUFR winds …")
    isfile(prepbufr_path) || error("PrepBUFR file not found at $prepbufr_path")

    @warn """
    ⚠️  MISSING NUMSCALE TABLE INTEGRATION

    Current implementation uses uniform sigma_multiplier=1.0 for all observations.
    Fortran GSI applies platform/pressure/condition-specific multipliers from fort.207
    (NUMSCALE table) that range from 0.5-5.0×, creating observation error heterogeneity.

    This causes:
    - Uniform observation errors → poor J_o/J_b balance
    - Weak conditioning → slow NAM CG convergence

    Workaround: Using increased background_wind_std and spatial thinning.
    TODO: Implement NUMSCALE table reader (Phase 2 of fix plan)
    """ maxlog=1

    backend_choice = bufr_backend
    use_bufrlib_flag = backend_choice == :eccodes ? false : true
    backend_choice ∈ (:auto, :bufrlib, :eccodes) || error("Unsupported bufr_backend option: $(backend_choice)")

    max_obs_request = max(200_000, max_count > 0 ? max(ceil(Int, max_count / 2) + 5_000, 100_000) : 200_000)
    wind_config = WindBUFRReader.WindBUFRConfig(
        use_bufrlib = use_bufrlib_flag,
        table_path = bufr_table_path,
        max_obs = max_obs_request,
        qc_threshold = 3,
        apply_10m_reduction = true,
        verbose = false
    )

    println(use_bufrlib_flag ? "  Decoder backend: BUFRLIB" : "  Decoder backend: ecCodes")

    bundle = WindBUFRReader.read_wind_bufr(prepbufr_path, wind_config)
    if bundle.nobs == 0
        error("Wind BUFR reader returned zero observations")
    end

    lon_axis = vec(lon_grid[1, :])
    lat_axis = vec(lat_grid[:, 1])

    values = Float64[]
    errors = Float64[]
    indices = Int[]
    components = Symbol[]
    platforms = Int[]

    dropped_domain = 0
    dropped_level = 0
    dropped_nan = 0
    dropped_quality = 0

    for obs in bundle.observations
        if !obs.use_obs
            dropped_quality += 1
            continue
        end

        if !(isfinite(obs.longitude) && isfinite(obs.latitude) && isfinite(obs.pressure))
            dropped_nan += 1
            continue
        end

        i = nearest_index_sorted(lon_axis, obs.longitude)
        j = nearest_index_sorted(lat_axis, obs.latitude)
        if i === nothing || j === nothing
            dropped_domain += 1
            continue
        end

        k = pressure_to_sigma_index(obs.pressure, sigma_axis)
        if k === nothing
            dropped_level += 1
            continue
        end

        platform_type = obs.platform_code
        sigma_mult = isfinite(obs.sigma_multiplier) && obs.sigma_multiplier > 0 ? obs.sigma_multiplier : 1.0
        dup_factor = obs.duplicate_factor > 0 ? obs.duplicate_factor : 1.0

        base_error = if use_platform_errors
            get(PLATFORM_ERROR_STDS, platform_type, DEFAULT_WIND_ERROR_STD)
        else
            DEFAULT_WIND_ERROR_STD
        end
        obs_error = max(error_floor, base_error * sigma_mult * sqrt(dup_factor))

        base_index = (k - 1) * nx * ny + (j - 1) * nx + i
        u_state_index = base_index
        v_state_index = nx * ny * nz + base_index

        if isfinite(obs.u)
            push!(values, obs.u)
            push!(errors, obs_error)
            push!(indices, u_state_index)
            push!(components, :u)
            push!(platforms, platform_type)
        end
        if isfinite(obs.v)
            push!(values, obs.v)
            push!(errors, obs_error)
            push!(indices, v_state_index)
            push!(components, :v)
            push!(platforms, platform_type)
        end
    end

    initial_count = length(values)

    dropped_thinning = 0
    if thinning_distance > 0.0 && !isempty(values)
        println(@sprintf("  Applying spatial thinning (%.1f km grid) ...", thinning_distance))
        grid_cfg = grid.config
        dx_km = grid_cfg.dx / 1000.0
        dy_km = grid_cfg.dy / 1000.0
        thin_nx = max(1, round(Int, thinning_distance / dx_km))
        thin_ny = max(1, round(Int, thinning_distance / dy_km))

        thinning_grid = Dict{Tuple{Int,Int,Int,Int}, Int}()
        keep_indices = Int[]
        for (obs_i, state_idx) in enumerate(indices)
            comp = components[obs_i]
            offset = comp === :u ? 0 : nx * ny * nz
            base = state_idx - offset

            k = div(base - 1, nx * ny) + 1
            remainder = (base - 1) % (nx * ny)
            j = div(remainder, nx) + 1
            i = remainder % nx + 1

            key = (div(i - 1, thin_nx), div(j - 1, thin_ny), k, comp === :u ? 0 : 1)
            if !haskey(thinning_grid, key)
                thinning_grid[key] = obs_i
                push!(keep_indices, obs_i)
            else
                dropped_thinning += 1
            end
        end

        values = values[keep_indices]
        errors = errors[keep_indices]
        indices = indices[keep_indices]
        components = components[keep_indices]
        platforms = platforms[keep_indices]

        println(@sprintf("  Thinning: %d → %d observations (dropped %d)",
                         initial_count, length(values), dropped_thinning))
    end

    dropped_maxcount = 0
    if max_count > 0 && length(values) > max_count
        println(@sprintf("  Applying max count limit (%d) ...", max_count))
        keep_indices = sort(shuffle(1:length(values))[1:max_count])
        dropped_maxcount = length(values) - max_count
        values = values[keep_indices]
        errors = errors[keep_indices]
        indices = indices[keep_indices]
        components = components[keep_indices]
        platforms = platforms[keep_indices]
        println(@sprintf("  Max count: %d → %d observations",
                         initial_count - dropped_thinning, length(values)))
    end

    counts = Dict(
        :total_records => bundle.nobs * 2,
        :u_records => bundle.nobs,
        :v_records => bundle.nobs,
        :assimilated => length(values),
        :dropped_domain => dropped_domain,
        :dropped_level => dropped_level,
        :dropped_nan => dropped_nan,
        :dropped_quality => dropped_quality,
        :dropped_thinning => dropped_thinning,
        :dropped_maxcount => dropped_maxcount,
        :decoder_status => Int(get(bundle.metadata, "decoder_status", 0))
    )

    println(@sprintf("  Winds available: raw=%d (u=%d, v=%d), assimilated=%d, dropped=%d",
                     bundle.nobs * 2, bundle.nobs, bundle.nobs, counts[:assimilated],
                     dropped_domain + dropped_level + dropped_nan + dropped_quality +
                     dropped_thinning + dropped_maxcount))

    return WindObservationPayload(values, errors, indices, components, platforms, counts)
end

function run_diagonal_analysis(background_state::Vector{Float64}, payload::WindObservationPayload,
                               nx::Int, ny::Int, nz::Int; background_wind_std::Float64=3.5)
    analysis_state = copy(background_state)
    increment = zeros(Float64, length(background_state))

    total_innov = 0.0
    total_weight = 0.0
    count = 0

    for (val, err, idx, comp) in zip(payload.values, payload.errors, payload.state_indices, payload.components)
        b_std = background_wind_std
        r = err^2
        b = b_std^2
        gain = b / (b + r)
        background = background_state[idx]
        increment_val = gain * (val - background)
        analysis_state[idx] += increment_val
        increment[idx] = increment_val
        total_innov += abs(val - background)
        total_weight += gain
        count += 1
    end

    residual_background = payload.values .- background_state[payload.state_indices]
    residual_analysis = payload.values .- analysis_state[payload.state_indices]
    weights = 1.0 ./ (payload.errors .^ 2)

    stats = Dict(
        :mode => "diag3dvar",
        :avg_innovation => count == 0 ? 0.0 : total_innov / count,
        :avg_gain => count == 0 ? 0.0 : total_weight / count,
        :J_observation_background => 0.5 * sum(weights .* (residual_background .^ 2)),
        :J_observation_analysis => 0.5 * sum(weights .* (residual_analysis .^ 2)),
        :normalized_residual_rms_background => sqrt(mean((residual_background .^ 2) .* weights)),
        :normalized_residual_rms_analysis => sqrt(mean((residual_analysis .^ 2) .* weights)),
        :analysis_increment_norm => norm(increment),
        :analysis_increment_rms => sqrt(mean(increment .^ 2))
    )

    return analysis_state, increment, stats
end

function create_nam_config(grid_cfg, output_dir; max_iter::Int=60, tol::Float64=5e-3, background_wind_std::Float64=3.5)
    AnalysisConfig{Float64}(
        grid_size = (grid_cfg.nx, grid_cfg.ny, grid_cfg.nsig),
        ensemble_size = 0,
        analysis_method = "3DVar",
        hybrid_coeff = 0.0,
        precision = Float64,
        use_hybrid = false,
        max_iterations = max_iter,
        convergence_tol = tol,
        model_type = "WRF",
        output_path = output_dir,
        params = Dict(
            "analysis_mode" => "NAM",
            "background_wind_std" => background_wind_std,
            "u_std" => background_wind_std,
            "v_std" => background_wind_std
        )
    )
end

function apply_binv!(dest::Vector{Float64}, ops::NAMOperators, src::Vector{Float64})
    copy!(ops.cv_workspace.values, src)
    result_cv = BackgroundError.apply_background_error_inverse(ops.B, ops.cv_workspace)
    copy!(dest, result_cv.values)
    return dest
end

function apply_HtRinvH!(dest::Vector{Float64}, vec::Vector{Float64},
                        payload::WindObservationPayload, weights::Vector{Float64})
    for idx in payload.state_indices
        dest[idx] = 0.0
    end
    @inbounds for (i, idx) in enumerate(payload.state_indices)
        dest[idx] += weights[i] * vec[idx]
    end
    return dest
end

function run_nam_analysis(background_state::Vector{Float64}, payload::WindObservationPayload,
                          grid_cfg, output_dir::AbstractString;
                          max_iter::Int=60, tol::Float64=5e-3, preconditioner::String="none",
                          use_balance::Bool=false, adaptive_tol::Bool=true, background_wind_std::Float64=3.5)
    config = create_nam_config(grid_cfg, output_dir; max_iter=max_iter, tol=tol, background_wind_std=background_wind_std)
    println("[4/6] Building operational B-operator …")
    B = BackgroundError.create_operational_bmatrix(config; use_balance=use_balance, use_vertical_modes=true)
    ops = NAMOperators(config, B, ControlVector(config))

    n = length(background_state)
    weights = 1.0 ./ (payload.errors .^ 2)
    residual_background = payload.values .- background_state[payload.state_indices]

    # Observation error statistics by component
    obs_error_stats = Dict{Symbol,Any}()
    for comp in [:u, :v]
        mask = payload.components .== comp
        if any(mask)
            errs = payload.errors[mask]
            obs_error_stats[comp] = Dict(
                :median => median(errs),
                :q25 => quantile(errs, 0.25),
                :q75 => quantile(errs, 0.75),
                :min => minimum(errs),
                :max => maximum(errs),
                :count => sum(mask)
            )
        end
    end

    # Add platform-wise statistics
    platform_counts = Dict{Int, Int}()
    for plat in payload.platforms
        platform_counts[plat] = get(platform_counts, plat, 0) + 1
    end
    obs_error_stats[:by_platform] = platform_counts

    # Show platform distribution
    println("  Observation platform distribution:")
    platform_names = Dict(
        120 => "RAOB", 130 => "AIRCRAFT", 180 => "SYNOP",
        181 => "SHIP", 182 => "BUOY", 183 => "BUOY", 187 => "METAR"
    )
    for (plat, count) in sort(collect(platform_counts), by=x->x[2], rev=true)
        name = get(platform_names, plat, "TYPE$plat")
        pct = 100.0 * count / length(payload.platforms)
        println(@sprintf("    %s (%-3d): %6d obs (%.1f%%)", name, plat, count, pct))
    end

    # Build preconditioner
    M_diag = ones(Float64, n)
    if preconditioner == "diag"
        println("  Applying diagonal preconditioner (Jacobi-style) …")
        # Build observation density map
        obs_density = zeros(Float64, n)
        for (i, idx) in enumerate(payload.state_indices)
            obs_density[idx] += weights[i]
        end
        # Preconditioner: M = diag(1 / sqrt(diag(A))) where A = H^T R^{-1} H + B^{-1}
        # Approximate: use observation density as H^T R^{-1} H contribution
        # Add small background contribution to avoid division by zero
        background_variance = 3.5^2  # m^2 s^{-2} for winds
        for i in 1:n
            if obs_density[i] > 0
                M_diag[i] = 1.0 / sqrt(obs_density[i] + 1.0 / background_variance)
            else
                M_diag[i] = sqrt(background_variance)
            end
        end
    elseif preconditioner == "block"
        @warn "Block preconditioner not yet implemented; falling back to none"
        preconditioner = "none"
    end

    # Apply preconditioning to rhs
    rhs = zeros(Float64, n)
    @inbounds for (i, idx) in enumerate(payload.state_indices)
        rhs[idx] += weights[i] * residual_background[i]
    end
    if preconditioner == "diag"
        rhs .*= M_diag
    end

    x = zeros(Float64, n)
    r = copy(rhs)
    p = copy(r)
    Ap = zeros(Float64, n)
    binv_p = zeros(Float64, n)
    obs_work = zeros(Float64, n)

    # CG iteration history
    cg_history = Dict{String,Any}[
        Dict(
            "iteration" => 0,
            "residual_norm" => norm(r),
            "alpha" => 0.0,
            "beta" => 0.0,
            "x_norm" => 0.0,
            "normalized_rms" => sqrt(mean((residual_background .^ 2) .* weights))
        )
    ]

    rs_old = dot(r, r)
    initial_residual_norm = sqrt(rs_old)  # For relative convergence test
    converged = false
    iterations = 0

    # Disable diagonal preconditioner due to implementation bug (double-application)
    if preconditioner == "diag"
        @warn "Diagonal preconditioner disabled due to implementation bug - using 'none' instead"
        preconditioner = "none"
    end

    if rs_old == 0.0
        converged = true
    else
        for iter in 1:max_iter
            iterations = iter
            apply_HtRinvH!(obs_work, p, payload, weights)
            if preconditioner == "diag"
                obs_work .*= M_diag
            end
            apply_binv!(binv_p, ops, p)
            if preconditioner == "diag"
                binv_p .*= M_diag
            end
            @. Ap = obs_work + binv_p
            denom = dot(p, Ap)
            if denom == 0.0
                @warn "Encountered zero denominator in CG solve; aborting early"
                break
            end
            alpha = rs_old / denom
            @. x = x + alpha * p
            @. r = r - alpha * Ap
            rs_new = dot(r, r)
            current_residual_norm = sqrt(rs_new)

            # Calculate current normalized RMS
            x_unpreconditioned = preconditioner == "diag" ? x ./ M_diag : x
            current_residual = payload.values .- (background_state[payload.state_indices] .+ x_unpreconditioned[payload.state_indices])
            current_normalized_rms = sqrt(mean((current_residual .^ 2) .* weights))

            push!(cg_history, Dict(
                "iteration" => iter,
                "residual_norm" => current_residual_norm,
                "alpha" => alpha,
                "beta" => 0.0,  # will update below
                "x_norm" => norm(x_unpreconditioned),
                "normalized_rms" => current_normalized_rms,
                "relative_residual" => current_residual_norm / initial_residual_norm
            ))

            # Convergence checks: use RELATIVE residual norm (matching Fortran GSI)
            relative_reduction = current_residual_norm / initial_residual_norm
            if relative_reduction < tol
                converged = true
                @info "CG converged (relative residual)" iter relative=relative_reduction tol=tol
                break
            end
            if adaptive_tol && current_normalized_rms < 1.2
                converged = true
                @info "CG converged (adaptive RMS criterion)" iter current_normalized_rms
                break
            end
            if iter % 10 == 0
                @info "CG progress" iter residual_norm=current_residual_norm relative_residual=relative_reduction normalized_rms=current_normalized_rms
            end

            beta = rs_new / rs_old
            cg_history[end]["beta"] = beta
            @. p = r + beta * p
            rs_old = rs_new
        end
    end

    # Undo preconditioning
    if preconditioner == "diag"
        x ./= M_diag
    end

    analysis_state = background_state + x
    analysis_residual = payload.values .- analysis_state[payload.state_indices]

    apply_binv!(binv_p, ops, x)
    Jb = 0.5 * dot(x, binv_p)
    Jo_background = 0.5 * sum(weights .* (residual_background .^ 2))
    Jo_analysis = 0.5 * sum(weights .* (analysis_residual .^ 2))

    stats = Dict(
        :mode => "nam",
        :iterations => iterations,
        :converged => converged,
        :preconditioner => preconditioner,
        :use_balance => use_balance,
        :J_background => Jb,
        :J_observation_background => Jo_background,
        :J_observation_analysis => Jo_analysis,
        :cost_reduction => Jo_background - (Jb + Jo_analysis),
        :normalized_residual_rms_background => sqrt(mean((residual_background .^ 2) .* weights)),
        :normalized_residual_rms_analysis => sqrt(mean((analysis_residual .^ 2) .* weights)),
        :analysis_increment_norm => norm(x),
        :analysis_increment_rms => sqrt(mean(x .^ 2)),
        :observation_error_stats => obs_error_stats,
        :cg_history => cg_history
    )

    return analysis_state, x, stats
end

function sample_radiances(radiance_files::Vector{Tuple{String,String}}; sample_limit::Int=10_000)
    if sample_limit <= 0
        println("[5/6] Skipping radiance sampling (limit ≤ 0)")
        return Dict{String,Any}()
    end

    println("[5/6] Sampling radiance BUFR files …")
    summary = Dict{String,Any}()
    config = RadianceBUFRConfig(max_obs=sample_limit, verbose=false)

    for (label, path) in radiance_files
        if !isfile(path)
            summary[label] = Dict("available" => false, "message" => "file not present")
            @warn "Radiance file missing" instrument=label path
            continue
        end
        try
            obs = read_radiance_bufr(path; config=config)
            n_obs = size(obs.brightness_temp, 2)
            summary[label] = Dict(
                "available" => true,
                "instrument" => obs.instrument,
                "channels" => length(obs.channels),
                "sampled_observations" => n_obs,
                "value_kind" => get(obs.metadata, "value_kind", "brightnessTemperature"),
            )
            println(@sprintf("  %s: instrument=%s, channels=%d, sampled=%d",
                             label, obs.instrument, length(obs.channels), n_obs))
        catch err
            summary[label] = Dict(
                "available" => false,
                "message" => "decode failed",
                "error" => sprint(showerror, err)
            )
            @warn "Radiance decode failed" label path exception=(err, catch_backtrace())
        end
    end

    return summary
end

function write_outputs(output_dir::AbstractString;
                       scenario::String,
                       analysis_mode::String,
                       background_summary::Dict,
                       counts::Dict{Symbol,Int},
                       assimilation_summary::Dict,
                       analysis_state::Vector{Float64},
                       background_state::Vector{Float64},
                       radiance_summary::Dict)
    mkpath(output_dir)

    metrics = Dict(
        "timestamp" => string(Dates.now()),
        "scenario" => scenario,
        "analysis_mode" => analysis_mode,
        "background" => Dict(string(k) => v for (k, v) in background_summary),
        "wind_observations" => Dict("counts" => Dict(string(k) => v for (k, v) in counts)),
        "analysis" => Dict(string(k) => v for (k, v) in assimilation_summary),
        "radiance_samples" => radiance_summary
    )

    metrics_path = joinpath(output_dir, "analysis_metrics.json")
    open(metrics_path, "w") do io
        JSON3.write(io, metrics; indent=2)
        println(io)
    end

    analysis_path = joinpath(output_dir, "analysis_state.bin")
    increment_path = joinpath(output_dir, "analysis_increment.bin")
    background_path = joinpath(output_dir, "background_state.bin")

    open(analysis_path, "w") do io
        write(io, reinterpret(UInt8, analysis_state))
    end
    open(background_path, "w") do io
        write(io, reinterpret(UInt8, background_state))
    end
    increment = analysis_state .- background_state
    open(increment_path, "w") do io
        write(io, reinterpret(UInt8, increment))
    end

    println(@sprintf("[6/6] Metrics written to %s", metrics_path))
    println(@sprintf("      Analysis state dump: %s", analysis_path))
    println(@sprintf("      Increment dump     : %s", increment_path))
end

function maybe_configure_prepbufr_table!(table_path::Union{Nothing,String})
    if table_path === nothing
        return
    end
    if !isfile(table_path)
        @warn "Specified PrepBUFR table file not found" table_path
        return
    end
    ENV["GSICORE_PREPBUFR_TABLE"] = table_path
    @info "Registered PrepBUFR table override" table_path
end

function main()
    (scenario, background_path, prepbufr_path, output_dir, radiance_sample, analysis_mode, bufr_table_path,
     bufr_backend, radiance_files, nam_max_iter, nam_tol, nam_preconditioner, nam_use_balance, nam_adaptive_tol,
     obs_thinning_distance, obs_max_count, obs_error_floor, obs_use_platform_errors, background_wind_std) = parse_cli_args()
    maybe_configure_prepbufr_table!(bufr_table_path)
    ensure_eccodes_overlay!()

    println("="^80)
    println("Julia Regional Case — $(uppercase(scenario)) configuration")
    println("="^80)
    println("Background file : $background_path")
    println("PrepBUFR winds  : $prepbufr_path")
    println("Output directory: $output_dir")
    if bufr_table_path !== nothing
        println("PrepBUFR table : $bufr_table_path")
    end
    println("BUFR backend  : $bufr_backend")
    println("Analysis mode   : $analysis_mode")
    if analysis_mode == "nam"
        println("NAM max_iter    : $nam_max_iter")
        println("NAM tolerance   : $nam_tol")
        println("NAM precond.    : $nam_preconditioner")
        println("NAM balance     : $nam_use_balance")
        println("NAM adaptive tol: $nam_adaptive_tol")
    end
    println("Obs thinning    : $(obs_thinning_distance > 0 ? "$(obs_thinning_distance) km" : "disabled")")
    println("Obs max count   : $(obs_max_count > 0 ? obs_max_count : "unlimited")")
    println("Obs error floor : $obs_error_floor m/s")
    println("Platform errors : $(obs_use_platform_errors ? "enabled" : "disabled")")
    println("Background σ_b  : $background_wind_std m/s")
    println("Scenario        : $scenario")
    println()

    bg = load_background(background_path)

    println("[2/6] Flattening background into state vector …")
    background_state = background_to_state_vector(bg)
    println(@sprintf("  State vector size: %d elements", length(background_state)))

    lon_grid, lat_grid, _, _, sigma_axis = get_coordinates(bg.grid)
    grid_cfg = bg.grid.config

    payload = build_wind_payload(prepbufr_path, bg.grid, background_state,
                                 lon_grid, lat_grid, sigma_axis,
                                 grid_cfg.nx, grid_cfg.ny, grid_cfg.nsig;
                                 bufr_table_path=bufr_table_path,
                                 bufr_backend=bufr_backend,
                                 thinning_distance=obs_thinning_distance,
                                 max_count=obs_max_count,
                                 error_floor=obs_error_floor,
                                 use_platform_errors=obs_use_platform_errors)

    if payload.counts[:assimilated] == 0
        error("No winds assimilated; cannot proceed with analysis")
    end

    analysis_state = Vector{Float64}()
    increment = Vector{Float64}()
    assimilation_stats = Dict{Symbol,Any}()

    if analysis_mode == "nam"
        analysis_state, increment, assimilation_stats = run_nam_analysis(
            background_state, payload, grid_cfg, output_dir;
            max_iter = nam_max_iter, tol = nam_tol,
            preconditioner = nam_preconditioner,
            use_balance = nam_use_balance,
            adaptive_tol = nam_adaptive_tol,
            background_wind_std = background_wind_std
        )
    else
        analysis_state, increment, assimilation_stats = run_diagonal_analysis(
            background_state, payload, grid_cfg.nx, grid_cfg.ny, grid_cfg.nsig;
            background_wind_std = background_wind_std
        )
    end

    radiance_summary = sample_radiances(radiance_files; sample_limit=radiance_sample)

    background_summary = Dict(
        "background_path" => get(bg.metadata, "source_file", background_path),
        "grid" => Dict(
            "nx" => grid_cfg.nx,
            "ny" => grid_cfg.ny,
            "nz" => grid_cfg.nsig,
            "dx_m" => grid_cfg.dx,
            "dy_m" => grid_cfg.dy,
        ),
        "state_size" => length(background_state),
        "state_mean" => mean(background_state),
        "state_std" => std(background_state)
    )

    write_outputs(output_dir;
                  scenario = scenario,
                  analysis_mode = analysis_mode,
                  background_summary=background_summary,
                  counts=payload.counts,
                  assimilation_summary=assimilation_stats,
                  analysis_state=analysis_state,
                  background_state=background_state,
                  radiance_summary=radiance_summary)

    # Save CG iteration history if available
    if haskey(assimilation_stats, :cg_history) && analysis_mode == "nam"
        timestamp = Dates.format(Dates.now(), "yyyymmdd")
        history_path = "prompt/proj7/logs/mwri_nam_cg_history_$(timestamp).json"
        mkpath(dirname(history_path))
        open(history_path, "w") do io
            JSON3.write(io, Dict(
                "timestamp" => string(Dates.now()),
                "scenario" => scenario,
                "config" => Dict(
                    "max_iter" => nam_max_iter,
                    "tol" => nam_tol,
                    "preconditioner" => nam_preconditioner,
                    "use_balance" => nam_use_balance,
                    "adaptive_tol" => nam_adaptive_tol
                ),
                "iterations" => assimilation_stats[:cg_history]
            ); indent=2)
            println(io)
        end
        println("CG iteration history written to $history_path")
    end

    println("Done. Assimilated $(payload.counts[:assimilated]) winds; metrics written to $output_dir")
end

main()
